{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "## Overview\n",
    "Here we use ASPECT to numerically reproduce the results of a linear stability analysis for the onset of convection in a fluid layer heated from below. This exercise was assigned to students in a geodynamics class at Portland State University as a first step towards setting up a nominally Earth-like mantle convection model. Hence, representative length scales and transport properties for Earth are used.\n",
    "\n",
    "The linear stability analysis appears in Turcotte and Schubert (2014) section 6.19. To use this code, you must compile aspect and give the path to the executable below as $\\texttt{aspect_bin}$. The critical Rayleigh number for the onset of convection depends only on the dimensionless wavelength of the perturbation, which is assumed to be equal to width of the domain. The domain has height $b$ and width $\\lambda$ and the perturbation has the functional form:\n",
    "\n",
    "$\n",
    "T'(x,y) = T_0'\\cos\\left(\\frac{2\\pi x}{\\lambda}\\right)\\sin\\left(\\frac{\\pi y}{b} \\right)\n",
    "$\n",
    "\n",
    "Note that because we place the bottom boundary of the domain at $y=0$ and the top at $y=b$, the perturbation vanishes at the top and bottom boundaries.\n",
    "\n",
    "The analytic solution for the critical Rayleigh number, $Ra_{cr}$ is given in T\\&S equation 6.319:\n",
    "\n",
    "$\n",
    "Ra_{cr}=\\frac{\\left(\\pi^2+\\frac{4\\pi^2 b^2}{\\lambda^2}\\right)^3}{\\frac{4\\pi^2 b^2}{\\lambda^2}}\n",
    "$\n",
    "\n",
    "The linear stability analysis also makes a prediction for the dimensionless growth rate of the instability $\\alpha'$. The maximum vertical velocity is given by:\n",
    "\n",
    "$\n",
    "v_{y,max} = \\frac{2\\pi}{\\lambda}\\phi_0' e^{\\alpha' t}\n",
    "$,\n",
    "\n",
    "where\n",
    "\n",
    "$\n",
    "\\phi_0' = -\\frac{2\\pi}{\\lambda}\\frac{\\rho_0 g \\alpha T_0'}{\\mu}\\left(\\frac{4\\pi^2}{\\lambda^2}+\\frac{\\pi^2}{b^2} \\right)^{-2}\n",
    "$,\n",
    "\n",
    "and\n",
    "\n",
    "$\n",
    "\\alpha'=\\frac{\\kappa}{b^2}\\left[\\frac{\\rho_0 g \\alpha b^3 \\Delta T}{\\mu \\kappa}\\left(\\frac{\\frac{4\\pi^2 b^2}{\\lambda^2}}{\\left(\\frac{4\\pi^2 b^2}{\\lambda^2}+\\pi^2\\right)^2}\\right) -\\left(\\pi^2+\\frac{4\\pi^2b^2}{\\lambda^2}\\right)\\right]\n",
    "$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "from subprocess import run\n",
    "\n",
    "base_input = \"convection-box-base.prm\"     # The 'base' input file that gets modified\n",
    "input_file = \"input.prm\"                   # The temporary input file that is used for each run\n",
    "aspect_bin = \"../../build/aspect\"          # Path to aspect executable\n",
    "output_dir = \"output\"                      # Output directory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "# Create a dictionary to replace keys in base input file with appropriate values\n",
    "parameters = dict([])\n",
    "parameters['PMAG'] = 1                # amplitude of temperature perturbation\n",
    "parameters['HEIGHT'] = 3.0e6          # layer thickness\n",
    "parameters['IREF'] = 1              # initial global refinement\n",
    "parameters['DELTA_T'] = 2500.0        # temperature difference between bottom, top boundaries\n",
    "parameters['GRAVITY'] = 10.0          # vertical gravitaitonal acceleration\n",
    "parameters['KTHERMAL'] = 4.0          # thermal conductivity\n",
    "parameters['ALPHA'] = 3e-5            # thermal expansivity\n",
    "parameters['DENSITY'] = 4000.0        # reference density\n",
    "parameters['SPECIFIC_HEAT'] = 1250.0  # specific heat capacity\n",
    "parameters['NY'] = 16;                # number of x-repetitions\n",
    "\n",
    "def generate_input_file(base_file_name,output_file_name,dictionary):\n",
    "    \"\"\"Read the 'base' input file from base_file_name, replace strings \n",
    "    using dictionary, and write new output file to output_file_name\"\"\"\n",
    "    fh = open(base_file_name,'r')\n",
    "    run(['rm','-f',output_file_name])\n",
    "    ofh = open(output_file_name,'w')\n",
    "    for line in fh:        \n",
    "        for key in dictionary:\n",
    "            if key in line:                \n",
    "                line = line.replace(key,str(dictionary[key]))\n",
    "        ofh.write(line)\n",
    "    fh.close()\n",
    "    ofh.close()\n",
    "    \n",
    "def parse_output(output_dir):\n",
    "    \"\"\"Read the statistics (stats_file) file from the output directory\"\"\"\n",
    "    tmp = np.loadtxt(output_dir + '/statistics',usecols=(1, 11))\n",
    "    vy = np.loadtxt(output_dir + '/point_values.txt',usecols=(4,))\n",
    "    t=tmp[:,0]\n",
    "    vmax = tmp[:,1]\n",
    "    return t,vmax,vy\n",
    "    \n",
    "def vy_exact(t,p):\n",
    "    \"\"\"Evaluate exact expression for maximum upward velocity\"\"\"\n",
    "    lam = p['WIDTH']\n",
    "    b = p['HEIGHT']\n",
    "    g = p['GRAVITY']\n",
    "    T0 = p['PMAG']\n",
    "    alpha = p['ALPHA']\n",
    "    eta = p['VISCOSITY']\n",
    "    rho = p['DENSITY']\n",
    "    dT = p['DELTA_T']\n",
    "    kappa = p['KTHERMAL']/rho/p['SPECIFIC_HEAT']\n",
    "    alpha_prime = kappa/b**2*((rho*g*alpha*b**3*dT/eta/kappa)*(4.0*np.pi**2*b**2/lam**2/(4.0*np.pi**2*b**2/lam**2+np.pi**2)**2)-(np.pi**2+4.0*np.pi**2*b**2/lam**2))\n",
    "    phi0 = -2.0*np.pi/lam*(rho*g*alpha*T0/eta)*(4.*np.pi**2/lam**2 + np.pi**2/b**2)**(-2.0)\n",
    "    vymax = 2.0*np.pi/lam*phi0 * np.exp(alpha_prime*t)\n",
    "    return vymax\n",
    "    \n",
    "def run_aspect(rayleigh,p=parameters):\n",
    "    \"\"\"Perofm the following tasks. \n",
    "    1. Choose viscosity to satisfy rayleigh number rayleigh and parameters in p, which\n",
    "    defaults to the parameters dictionary defined in the current namespace. \n",
    "    2. Generate an input file. \n",
    "    3. Run aspect\n",
    "    4. Parse the results, return the rate at which velocities increase between the first and second timesteps.\"\"\"\n",
    "    def eta(ra):        \n",
    "        rho = p['DENSITY']\n",
    "        alpha = p['ALPHA']\n",
    "        g = p['GRAVITY']\n",
    "        dT = p['DELTA_T']\n",
    "        h = p['HEIGHT']     \n",
    "        kappa = p['KTHERMAL']/p['DENSITY']/p['SPECIFIC_HEAT']    \n",
    "        return rho*g*alpha*dT*(h**3)/ra/kappa                \n",
    "    \n",
    "    p['VISCOSITY'] = eta(rayleigh);\n",
    "    # Change the number of x-repeitions to keep elements close to equant\n",
    "    p['NX'] = int(p['WIDTH']/p['HEIGHT']*p['NY'])\n",
    "    p['EVALUATION_POINTS'] =  str(0.0) + ',' + str(p['HEIGHT']/2.0)\n",
    "    generate_input_file(base_input,input_file,p)\n",
    "    run(['rm','-rf',output_dir])\n",
    "    run([aspect_bin,input_file])\n",
    "    t, vmax, vy = parse_output(output_dir)\n",
    "    return (vmax[1]-vmax[0]), t, vy\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "Define a function to perform bisection: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": true,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "def bisection(function, low, high, atol=np.Inf, rtol=0.01):\n",
    "    \"\"\"Performs bisection to find 0.0=function(x) using function and initial bounds (low, high) \n",
    "    to within absolute value tolerance atol OR relative tolerance rtol\"\"\"\n",
    "    assert function(low)[0]<0 \n",
    "    assert function(high)[0]>0\n",
    "    x_try = [];\n",
    "    y_try = [];\n",
    "    \n",
    "    while abs(low-high)>atol or abs(low-high)/abs(low) > rtol:\n",
    "        x = (low + high)/2.0\n",
    "        x_try.append(x)\n",
    "        y = function(x)[0]       \n",
    "        y_try.append(y)\n",
    "        if y > 0:\n",
    "            high=x;\n",
    "        else:\n",
    "            low=x;        \n",
    "    return x, x_try, y_try"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "deletable": true,
    "editable": true
   },
   "source": [
    "Define the range of aspect ratios for which $Ra_{cr}$ is calculated. Iterate over them and perform bisection to determine $Ra_{cr}$ for each."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "aspect_min = np.pi/4.\n",
    "aspect_max = 12.*np.pi\n",
    "naspect = 10;\n",
    "ra_rtol = 1e-5;\n",
    "\n",
    "def racr_lsa( b_over_lam ):\n",
    "    \"\"\"Evaluates the critical Rayleigh number for b_over_lam, which is the ratio of box height \n",
    "    to box width. Implements Turcotte and Schubert (3ed) equation 6.319\"\"\"\n",
    "    pi = np.pi\n",
    "    racr = (pi**2 + 4*pi**2*b_over_lam**2)**3/(4*pi**2*b_over_lam**2)\n",
    "    return racr\n",
    "\n",
    "aspect_ratio = np.linspace(1.0/aspect_max,1.0/aspect_min,naspect)\n",
    "aspect_ratio = 1.0/aspect_ratio\n",
    "lams=[]\n",
    "ra_save = []\n",
    "val_save = []\n",
    "racr_save = []\n",
    "xplot = [];\n",
    "for a in aspect_ratio:\n",
    "    # calculate dimensional width\n",
    "    parameters['WIDTH'] = a*parameters['HEIGHT']\n",
    "    racr_analytic =racr_lsa( 1.0/a )\n",
    "    racr_low  = racr_analytic/3.0\n",
    "    racr_high = racr_analytic*3.0\n",
    "    racr, ratry, vals = bisection(run_aspect,racr_low,racr_high,rtol=ra_rtol)\n",
    "    xplot.append( 2.0*np.pi*parameters['HEIGHT']/parameters['WIDTH'] )\n",
    "    ra_save.append(ratry)\n",
    "    racr_save.append(racr)\n",
    "    val_save.append(vals)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "print(parameters['WIDTH']/parameters['HEIGHT'])\n",
    "print(run_aspect(racr_high))\n",
    "print(run_aspect(racr_low))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true,
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def symbol(value):\n",
    "    \"\"\"Returns a string specifying figure glyph and color corresponding to value\"\"\"\n",
    "    if value>0:\n",
    "        return 'ro'\n",
    "    else:\n",
    "        return 'k.'\n",
    "\n",
    "plt.figure()\n",
    "for j in range(len(ra_save)):\n",
    "    # Plot the results\n",
    "    ratry = ra_save[j]\n",
    "    dvdt = val_save[j]\n",
    "    \n",
    "    for i in range(len(ratry)):\n",
    "        shape, mycolor = symbol(ratry[i])\n",
    "        plt.plot(xplot[j],ratry[i],symbol(dvdt[i]))\n",
    "\n",
    "# Plot the analytic solution.\n",
    "xp = np.linspace(xplot[0],xplot[-1],1000)\n",
    "plt.plot(xp,[racr_lsa(x/(2*np.pi)) for x in xp],'g:')\n",
    "plt.text(2,3000,'Unstable')\n",
    "plt.text(5.5,1000,'Stable')\n",
    "plt.xlabel('$2\\pi b/\\lambda$')\n",
    "plt.ylabel('$Ra_{cr}$')\n",
    "plt.ylim([0, 4500])\n",
    "\n",
    "plt.savefig('racr.png', bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "# Make a plot showing the relative error between numerical and analytic solutions\n",
    "# Recall that we should expect the level agreement to be limited by the relative error\n",
    "# tolerance defined during the bisection procedure.\n",
    "plt.figure()\n",
    "racr = np.array([racr_lsa(x) for x in xplot])\n",
    "e_norm = [np.abs(racr_save[i]-racr_lsa(xplot[i]/(2.*np.pi)))/np.abs(racr_lsa(xplot[i]/(2.*np.pi))) for i in range(len(xplot))]\n",
    "plt.plot(xplot,e_norm,'k.')\n",
    "plt.plot([0,8],[ra_rtol,ra_rtol],'r--')\n",
    "plt.xlabel('$2\\pi b/\\lambda$')\n",
    "plt.ylabel('relative error')\n",
    "plt.yscale('log')\n",
    "plt.savefig('racr_error.png', bbox_inches='tight')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "collapsed": false,
    "deletable": true,
    "editable": true
   },
   "outputs": [],
   "source": [
    "# As an additional example, compare the analytic and numerically determined maximal velocities\n",
    "parameters['WIDTH'] = 2.0*np.sqrt(2.0)*parameters['HEIGHT'];\n",
    "ra_values = np.linspace(250,1000,10)\n",
    "dvdt=[]\n",
    "vy0=[];  #Numerically determined vy at time 0\n",
    "vy1=[];  #Numerically determined vy at time 1\n",
    "vye0=[]; #Exact vy at time 0\n",
    "vye1=[]; #Exact vy at time 1\n",
    "for ra in ra_values:    \n",
    "    tmp1,tmp2,tmp3 = run_aspect(ra,parameters)\n",
    "    t0=tmp2[0]\n",
    "    t1=tmp2[1]\n",
    "    vye0.append( vy_exact(t0,parameters) )\n",
    "    vye1.append( vy_exact(t1,parameters) )\n",
    "    dvdt.append(tmp1)\n",
    "    vy0.append(tmp3[0])\n",
    "    vy1.append(tmp3[1])\n",
    "    \n",
    "err0 = [abs(vy0[i]-vye0[i])/abs(vye0[i]) for i in range(len(vy0))]\n",
    "err1 = [abs(vy1[i]-vye1[i])/abs(vye1[i]) for i in range(len(vy1))]\n",
    "plt.figure()\n",
    "plt.plot(ra_values,err0,'rx',label='Time 0')\n",
    "plt.plot(ra_values,err1,'kx',label='Time 1')\n",
    "plt.yscale('log')\n",
    "plt.legend()\n",
    "plt.ylabel('Relative Error')\n",
    "plt.xlabel('Rayleigh Number')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
